# whisper_asr.py
import os
import re
import opencc
from faster_whisper import WhisperModel

class WhisperASR:
    def __init__(self, model_path: str = None, device: str = "cuda", compute_type: str = "float16"):
        os.environ["HF_HUB_OFFLINE"] = "1"
        os.environ["TRANSFORMERS_OFFLINE"] = "1"
        os.environ["HF_DATASETS_OFFLINE"] = "1"

        self.converter = opencc.OpenCC('t2s')
        self.model = WhisperModel(
            "medium",
            device=device,
            compute_type=compute_type,
            download_root=model_path
        )

    def transcribe(self, audio_path: str):
        segments, info = self.model.transcribe(audio_path, beam_size=5, task="transcribe")
        results = []
        total_duration = 0.0
        total_chars = 0
        total_words = 0

        for seg in segments:
            text = self.converter.convert(seg.text) if info.language == "zh" else seg.text
            duration = seg.end - seg.start
            total_duration += duration

            # 中文字符数（不含标点）
            if info.language == "zh":
                char_count = len(re.findall(r'[\u4e00-\u9fff]', text))
                total_chars += char_count
            else:
                word_count = len(text.split())
                total_words += word_count

            results.append({
                "start": seg.start,
                "end": seg.end,
                "text": text
            })

        # 计算整体语速
        if info.language == "zh":
            overall_speed = total_chars / total_duration if total_duration > 0 else 0
            speed_unit = "字/秒"
        else:
            overall_speed = total_words / total_duration if total_duration > 0 else 0
            speed_unit = "words/s"

        return {
            "language": info.language,
            "segments": results,
            "overall_speed": round(overall_speed, 2),
            "speed_unit": speed_unit
        }